#include <iostream>
#include <chrono>
#include <thread>
#include <torch/script.h>

using namespace std;

void th_function1()  
{  
	torch::jit::script::Module sci_net_;
	sci_net_ = torch::jit::load("/home/st/ubuntu_data/NRSL/deep_learning/pytorch_learning/output/cnn_trace1.pt");
	torch::Tensor sci_tensor = torch::randn({1,1,28,28}, at::device(at::kCUDA).dtype(at::kFloat));
	torch::Tensor output;
	for(int i = 0; i<50000; i++){
		output = sci_net_.forward({sci_tensor}).toTensor();
	}
    std::cout << "hello thread1.  " << output << std::endl;  
}  

void th_function2()  
{  
	torch::jit::script::Module sci_net_;
	sci_net_ = torch::jit::load("/home/st/ubuntu_data/NRSL/deep_learning/pytorch_learning/output/cnn_trace2.pt");
	torch::Tensor sci_tensor = torch::randn({1,1,28,28}, at::device(at::kCUDA).dtype(at::kFloat));
	torch::Tensor output;
	for(int i = 0; i<50000; i++){
		output = sci_net_.forward({sci_tensor}).toTensor();
	}
    std::cout << "hello thread2.  " << output << std::endl;  
}  

int main(){

	std::chrono::system_clock::time_point t1, t2;
	t1 = std::chrono::system_clock::now();
	// th_function1();
	// th_function2();
	thread th1(th_function1);
	thread th2(th_function2);
	th1.join();
	th2.join();
	t2 = std::chrono::system_clock::now();
	cout<<std::chrono::duration_cast<std::chrono::microseconds>( t2-t1 ).count()<<std::endl;

	return 0;
}
